import numpy as np
import os
import output_function as of
import pandas as pd
import matplotlib.pyplot as plt

# 使用示例
base_path = './data/RB_Ra=4000_Lx=5.0/'  # 根据需要修改为您的保存路径
out_iter =40000  # 假设我们要加载第300时间步的数据

# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')

n_density=n

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
Ra=Ra_star = config['Ra_star']
prandtl=Pr = config['Pr']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

print(Ny)

# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
# 初始化n_ave数组，用于存储每个j值对应的n的平均值
n_ave = np.zeros(Ny)  # Ny为n数组在y轴方向的长度
n_init_ave=np.zeros(Ny)
radial_x=np.zeros(Ny)
# 对于每个j值，计算对应列的平均值
for j in range(Ny):
    n_ave[j] = np.mean(n[j, :])
    n_init_ave[j]=np.mean(n_init0_matrix[j, :])
    radial_x[j]=y[j]

# 绘制n_ave与y的关系图
plt.figure(figsize=(10, 6))
plt.plot(radial_x, n_init_ave, label='Initial', color='blue')
plt.plot(radial_x, n_ave, label=f'Ra={Ra_star}', color='red')

base_path = './data/RB_Ra=20000_Lx=5.0/'  # 根据需要修改为您的保存路径
out_iter =30000  # 假设我们要加载第300时间步的数据
# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')
config = of.load_simulation_config(f'{base_path}simulation_config.json')
Ny = config['Ny']
Ra_star = config['Ra_star']

n_ave = np.zeros(Ny)  # Ny为n数组在y轴方向的长度
n_init_ave=np.zeros(Ny)
radial_x=np.zeros(Ny)
# 对于每个j值，计算对应列的平均值
for j in range(Ny):
    n_ave[j] = np.mean(n[j, :])
    n_init_ave[j]=np.mean(n_init0_matrix[j, :])
    radial_x[j]=y[j]
plt.plot(radial_x, n_ave, label=f'Ra={Ra_star}', color='purple')

base_path = './data/RB_Ra=80000_Lx=5.0/'  # 根据需要修改为您的保存路径
out_iter =20000  # 假设我们要加载第300时间步的数据
# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')
config = of.load_simulation_config(f'{base_path}simulation_config.json')
Ny = config['Ny']
Ra_star = config['Ra_star']

n_ave = np.zeros(Ny)  # Ny为n数组在y轴方向的长度
n_init_ave=np.zeros(Ny)
radial_x=np.zeros(Ny)
# 对于每个j值，计算对应列的平均值
for j in range(Ny):
    n_ave[j] = np.mean(n[j, :])
    n_init_ave[j]=np.mean(n_init0_matrix[j, :])
    radial_x[j]=y[j]
plt.plot(radial_x, n_ave, label=f'Ra={Ra_star}', color='green')

plt.axvline(x=Ly/2, color='black', linestyle='--', linewidth=1, label='y midplane')
plt.xlabel('y')
plt.ylabel(r'Average $<T>_x$ ')
plt.title(f'Steady T profile of different Rayleigh number')
plt.legend()
plt.grid(True)
plt.show()
exit(0)







